# 导入包
from matplotlib import pyplot as plt
import pandas as pd
import statsmodels.formula.api as smf
from sqlalchemy import create_engine
import pymysql
import statsmodels.api as sm
import numpy as np
import seaborn as sns


db_config = {
    'host': '127.0.0.1',
    'user': 'root',
    'password': 'root',
    'database': 'tushare',
    'port': 3306,
    'charset': 'utf8mb4'   # 添加字符集设置
}
engine = create_engine(
    f"mysql+pymysql://{db_config['user']}:{db_config['password']}@{db_config['host']}:{db_config['port']}/{db_config['database']}?charset={db_config['charset']}"
)
conn = pymysql.connect(**db_config)
chunk_size = 10000
df = pd.read_sql_query("""
    SELECT d.* FROM date_1 d
    join moneyflows m on m.ts_code=d.ts_code and d.trade_date=m.trade_date
    left join undex_daily i on i
    WHERE d.trade_date BETWEEN '2023-01-01' and '2023-12-31' and d.ts_code = '000001.SZ'
    """,
    conn,
    chunksize=chunk_size)
df1 = pd.concat(df, ignore_index=True)

df1['zd_closes']=((df1['closes'] - df1['closes'].shift(1))/df1['closes'].shift(1),2)
df1['zd_closes']=((df1['i_vloses'] - df1['i_vloses'].shift(1))/df1['i_vloses'].shift(1),2)
df1['zd_vlo']=((df1['i_vol'] - df1['i_vol'].shift(1))/df1['i_vol'].shift(1),2)
print (df1.head)

df1=df1.dropna(subset=['zd_clases','zs_closes','zs_closes'])
print(df1.hesd)

ex=['zd_clases','id','ts_code','trade_date','the_date','opens','hogh','low','closes','pre_closes','changes','pct_chg','amount']
number=df1.select_dtypes(include=['number']).columns.tolist()

newList=[col for col in number if col not in ex]

format = 'zd_closes ~'+'+'.join(newList)
res=smf.ols(formuls,sata=df1).fit()
print(res.summary())

plt.figure(figsize=(10,6))
plt.scatter(res.fittedvalues,df1['zd_closes'])
plt.show()

features=df1[['buy_lg_vol','aell_lg_vol','sell_elg_vol','net_mf_vol','zs_closes']]
correlation_matrix=features.corr()

pit.figure(figsise=(10,6))
sns.heatmap(correlation_matrix,annot=True,cmap='coolwarm',cbar=True,fmt='.2f',linewidths=.5)
plt.show()